# Ollama 模型初始化脚本
# Windows PowerShell 7.5 兼容版本

param(
    [string]$ModelName = "",
    [switch]$Verbose
)

# 设置错误处理
$ErrorActionPreference = "Stop"
$VerbosePreference = if ($Verbose) { "Continue" } else { "SilentlyContinue" }

# 颜色输出函数
function Write-ColorOutput {
    param(
        [string]$Message,
        [string]$Color = "White"
    )
    Write-Host $Message -ForegroundColor $Color
}

function Write-Success {
    param([string]$Message)
    Write-ColorOutput "✅ $Message" "Green"
}

function Write-Error {
    param([string]$Message)
    Write-ColorOutput "❌ $Message" "Red"
}

function Write-Warning {
    param([string]$Message)
    Write-ColorOutput "⚠️  $Message" "Yellow"
}

function Write-Info {
    param([string]$Message)
    Write-ColorOutput "ℹ️  $Message" "Cyan"
}

# 获取脚本所在目录
$ScriptDir = Split-Path -Parent $MyInvocation.MyCommand.Path
Set-Location $ScriptDir

Write-ColorOutput "🚀 初始化 Ollama 模型..." "Magenta"

# 加载环境变量
$envFile = Join-Path $ScriptDir ".env"
if (Test-Path $envFile) {
    Write-Info "加载环境配置..."
    Get-Content $envFile | Where-Object { $_ -match '^[^#].*=' } | ForEach-Object {
        $key, $value = $_.Split('=', 2)
        $key = $key.Trim()
        $value = $value.Trim()
        [Environment]::SetEnvironmentVariable($key, $value)
        Write-Verbose "设置环境变量: $key = $value"
    }
}

# 确定要拉取的模型
if (-not $ModelName) {
    $ModelName = $env:OLLAMA_MODEL_NAME
}

if (-not $ModelName) {
    Write-Warning "未指定模型名称，使用默认值: qwen2.5:7b"
    $ModelName = "qwen2.5:7b"
}

Write-Info "目标模型: $ModelName"

# 检查 Ollama 是否安装
Write-Info "检查 Ollama 安装..."
try {
    $ollamaVersion = & ollama --version 2>$null
    if ($ollamaVersion) {
        Write-Success "Ollama 已安装: $ollamaVersion"
    } else {
        throw "Ollama 命令无响应"
    }
} catch {
    Write-Error "未检测到 Ollama 安装: $_"
    Write-Info "请从 https://ollama.com/download 安装 Ollama"
    exit 1
}

# 检查 Ollama 服务是否运行
Write-Info "检查 Ollama 服务状态..."
$ollamaBaseUrl = $env:OLLAMA_BASE_URL
if (-not $ollamaBaseUrl) {
    $ollamaBaseUrl = "http://localhost:11434"
}

try {
    $response = Invoke-WebRequest -Uri "$ollamaBaseUrl/api/tags" -TimeoutSec 5 -ErrorAction SilentlyContinue
    if ($response.StatusCode -eq 200) {
        Write-Success "Ollama 服务正在运行"
    } else {
        Write-Warning "Ollama 服务响应异常，尝试启动..."
        Start-Process -FilePath "ollama" -ArgumentList "serve" -NoNewWindow
        Start-Sleep -Seconds 3
    }
} catch {
    Write-Warning "Ollama 服务未运行，尝试启动..."
    try {
        Start-Process -FilePath "ollama" -ArgumentList "serve" -NoNewWindow
        Start-Sleep -Seconds 3
        Write-Success "Ollama 服务启动成功"
    } catch {
        Write-Error "无法启动 Ollama 服务: $_"
        exit 1
    }
}

# 检查模型是否已安装
Write-Info "检查模型 $ModelName 是否已安装..."
try {
    $installedModels = & ollama list 2>$null
    $modelExists = $false

    if ($installedModels) {
        foreach ($line in ($installedModels -split "`n")) {
            if ($line -match $ModelName) {
                $modelExists = $true
                break
            }
        }
    }

    if ($modelExists) {
        Write-Success "模型 $ModelName 已存在，跳过拉取"
    } else {
        Write-Info "模型 $ModelName 未安装，开始拉取..."
        Write-Info "这可能需要一些时间，请耐心等待..."

        $pullProcess = Start-Process -FilePath "ollama" -ArgumentList "pull", $ModelName -NoNewWindow -PassThru

        # 显示进度
        $spinner = @('|', '/', '-', '\')
        $spinnerIndex = 0
        while (-not $pullProcess.HasExited) {
            Write-Host "`r$($spinner[$spinnerIndex]) 正在拉取模型..." -NoNewline
            $spinnerIndex = ($spinnerIndex + 1) % $spinner.Length
            Start-Sleep -Milliseconds 500
        }
        Write-Host "`r" -NoNewline

        # 检查拉取结果
        if ($pullProcess.ExitCode -eq 0) {
            Write-Success "模型 $ModelName 拉取完成"
        } else {
            Write-Error "模型拉取失败 (退出码: $($pullProcess.ExitCode))"
            exit 1
        }
    }
} catch {
    Write-Error "检查/拉取模型失败: $_"
    exit 1
}

# 验证模型安装
Write-Info "验证模型安装..."
try {
    $installedModels = & ollama list 2>$null
    $modelFound = $false

    if ($installedModels) {
        Write-Info "当前已安装的模型:"
        foreach ($line in ($installedModels -split "`n")) {
            if ($line -match $ModelName) {
                $modelFound = $true
                Write-ColorOutput "  ✅ $line" "Green"
            } elseif ($line.Trim()) {
                Write-ColorOutput "     $line" "Gray"
            }
        }
    }

    if ($modelFound) {
        Write-Success "模型 $ModelName 验证通过"
    } else {
        Write-Error "模型 $ModelName 验证失败"
        exit 1
    }
} catch {
    Write-Warning "无法验证模型安装: $_"
}

# 显示使用提示
Write-ColorOutput "" "Magenta"
Write-Success "模型初始化完成！"
Write-ColorOutput "" "Magenta"
Write-Info "📦 已安装的模型: $ModelName"
Write-Info "🚀 现在可以启动 AI Agent 服务:"
Write-Info "   .\start.ps1"
Write-ColorOutput "" "Magenta"
Write-Info "💡 其他可用操作:"
Write-Info "   - 测试配置: .\test-ollama.ps1"
Write-Info "   - 查看模型: ollama list"
Write-Info "   - 删除模型: ollama rm $ModelName"
